#include "library.h"

// 全局变量
WNDPROC g_originWndProc = nullptr;      // 游戏本身的窗口处理函数

// 自己的窗口处理函数
LRESULT CALLBACK SelfWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam);

// ImGui窗口处理函数
extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

// Hooks
InlineHook* g_resetHook = nullptr;
InlineHook* g_presentHook = nullptr;

// Self handler
HRESULT WINAPI Hook_Reset (IDirect3DDevice9* direct3DDevice9, D3DPRESENT_PARAMETERS* pPresentationParameters);
HRESULT WINAPI Hook_Present(IDirect3DDevice9* direct3DDevice9, RECT* pSourceRect, RECT* pDestRect, HWND hDestWindowOverride, RGNDATA* pDirtyRegion);

// 初始化D3D9并hook函数
BOOL WINAPI InitializeD3D9() {

    // 初始化数据
    InitializeData();

    // 创建D3D
    IDirect3D9* direct3D9 = Direct3DCreate9(D3D_SDK_VERSION);
    if (direct3D9 == nullptr) {

        MessageBoxA(nullptr, "Direct3DCreate9 failed.", "Error:", MB_OK | MB_ICONERROR);
        return FALSE;
    }

    D3DPRESENT_PARAMETERS g_info;
    memset(&g_info, 0, sizeof(g_info));
    g_info.Windowed = true;
    g_info.SwapEffect = D3DSWAPEFFECT_COPY;

    // 创建设备，为了拿函数地址并Hook
    IDirect3DDevice9* direct3DDevice9 = nullptr;
    direct3D9->CreateDevice(D3DADAPTER_DEFAULT, D3DDEVTYPE_HAL, g_gameHwnd, D3DCREATE_SOFTWARE_VERTEXPROCESSING, &g_info, &direct3DDevice9);
    if (direct3DDevice9 == nullptr) {

        MessageBoxA(nullptr, "CreateDevice failed.", "Error:", MB_OK | MB_ICONERROR);
        return FALSE;
    }

    // Hook Direct3D
    int* direct3DDeviceTable = (int*)*(int*) direct3DDevice9;
    g_resetHook = new InlineHook(direct3DDeviceTable[16], reinterpret_cast<DWORD> (Hook_Reset));
    g_presentHook = new InlineHook(direct3DDeviceTable[17], reinterpret_cast<DWORD> (Hook_Present));
    g_resetHook->MotifyASM();
    g_presentHook->MotifyASM();

    // 销毁设备，因为拿到地址了
    direct3D9->Release();
    direct3DDevice9->Release();

    return TRUE;
}

// Hooked function implements
HRESULT WINAPI Hook_Reset (IDirect3DDevice9* direct3DDevice9, D3DPRESENT_PARAMETERS* pPresentationParameters) {

    g_resetHook->ResetASM();

    // 废除设备对象
    ImGui_ImplDX9_InvalidateDeviceObjects();

    HRESULT res = direct3DDevice9->Reset(pPresentationParameters);

    // 重新创建设备对象
    ImGui_ImplDX9_CreateDeviceObjects();

    g_resetHook->MotifyASM();

    return res;
}

HRESULT WINAPI Hook_Present(IDirect3DDevice9* direct3DDevice9, RECT* pSourceRect, RECT* pDestRect, HWND hDestWindowOverride, RGNDATA* pDirtyRegion) {

    g_presentHook->ResetASM();

    static bool called = false;
    if(!called) {

        called = true;

        // 初始化ImGui
        InitializeImGui(g_gameHwnd, direct3DDevice9);

        // 绑定热键
        BindHotKeys();

        // 接管窗口消息
        g_originWndProc = reinterpret_cast<WNDPROC> (SetWindowLongA(g_gameHwnd, GWL_WNDPROC, reinterpret_cast<long> (SelfWndProc)));
    }

    // 绘制
    ImGui_ImplDX9_NewFrame();
    ImGui_ImplWin32_NewFrame();
    ImGui::NewFrame();

    UpdateData();
    []() -> void {

      TraceRay::OnRender2D(hEngineDll, hClientDll);
    }();
    onRenderGui();

    ImGui::EndFrame();
    ImGui::Render();
    ImGui_ImplDX9_RenderDrawData(ImGui::GetDrawData());

    HRESULT res = direct3DDevice9->Present(pSourceRect, pDestRect, hDestWindowOverride, pDirtyRegion);

    g_presentHook->MotifyASM();

    // 调用原来的
    return res;
}

LRESULT CALLBACK SelfWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {

    // 更新热键状态
    HotKeyManager::KeyboardHandler(hWnd, uMsg, wParam, lParam);

    // ImGui接管窗口处理
    if (ImGui_ImplWin32_WndProcHandler(hWnd, uMsg, wParam, lParam)) return true;

    // 调用原先的窗口处理
    return CallWindowProcA(g_originWndProc, hWnd, uMsg, wParam, lParam);
}
